package com.example.chat.agent.service;

import dev.langchain4j.model.zhipu.ZhipuAiEmbeddingModel;
import dev.langchain4j.store.embedding.chroma.ChromaEmbeddingStore;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;

import java.time.Duration;

@Service
public class ChromaDbService {
    @Value("${zhipu.api.key}")
    private String apiKey;

    @Value("${chromadb.base-url}")
    private String dbBaseUrl;

    public ChromaEmbeddingStore getEmbeddingStore(String collectionName) {
        return ChromaEmbeddingStore
                .builder()
                .baseUrl(dbBaseUrl)
                .collectionName(collectionName)
                .logRequests(true)
                .logResponses(true)
                .timeout(Duration.ofSeconds(3))
                .build();
    }

    public ZhipuAiEmbeddingModel getEmbeddingModel() {
        return ZhipuAiEmbeddingModel.builder()
                .apiKey(apiKey)
                .logRequests(true)
                .logResponses(true)
                .maxRetries(3)
                .build();
    }
}
